import networkx as nx

import types
import copy
import inspect
from itertools import izip

from core import util, cmds

class Edge(object):
    def __init__(self, out_node, out_attr, in_node, in_attr):
        self.out_node = out_node
        self.out_attr = out_attr
        self.in_node = in_node
        self.in_attr = in_attr
        
    def out_data(self):
        return getattr(self.out_node, self.out_attr)
    
    def in_data(self):
        return getattr(self.in_node, self.in_attr)
    
    def push(self):
        out_data = copy.deepcopy(self.out_data())
        print 'push', out_data, 'to', edge.in_node, edge.in_attr
        self.in_node.set_data(self.in_attr, data)


class Pipeline(ni.Component):
    
    def init(self):
        self.graph = nx.MultiDiGraph()
    '''
    
    def __save__(self):
        data = {'nodes':self.nodes()}
        data['edges'] = []
        for out_node, in_node, d in self.graph.edges(data=True):
            data['edges'].append([out_node, d['out_attr'], in_node, d['in_attr']])
        return data
    '''
        
    def __load__(self, data):
        for node in data.get('nodes', []):
            self.graph.add_node(node, dirty=True)
        
        for key, edge in data.get('edges',{}):
            self.graph.add_edge(edge['out_node'], edge['in_node'], key=key,
                                in_attr=edge['in_attr'], out_attr=edge['out_attr']) 
    
    def add_node(self, node):
        self.graph.add_node(node, dirty=True)
        if not self._custom_data.has_key('nodes'):
            self._custom_data['nodes'] = []
        self._custom_data['nodes'].append(node)
        self._custom_data._p_note_change()
        cmds.save()
    
    def add_nodes(self, *nodes):
        map(self.add_node, nodes)
        
    def nodes(self, data=None):
        return self.graph.nodes(data)
    
    def edges(self, attrs=None):
        for out_node, in_node, data in self.graph.edges(data=True):
            if attrs is None:
                yield (out_node, in_node)
                continue
            out_attr = out_node.get_attr(data['out_attr'])
            in_attr = in_node.get_attr(data['in_attr'])
            yield (out_node, out_attr, in_node, in_attr)
        
    def remove_node(self, node):
        self.graph.remove_node(node)
        self._custom_data['nodes'].remove(node)
        self._custom_data._p_note_change()
        cmds.save()
        
    def remove_edge(self, out_node, out_attr_name, in_node, in_attr_name):
        key = self.get_key(out_attr_name, in_attr_name)
        self.graph.remove_edge(out_node, in_node, key=key)
        del self._custom_data['edges'][key]
        self._custom_data._p_note_change()
        cmds.save()
        
    
        
    def has_connection_to(self, node, in_attr_name):
        for out_node, in_node, data in self.graph.edges(data=True):
            if in_node is node and data['in_attr'] == in_attr_name:
                return True
        
    def connect(self, out_node, out_attr_name, in_node, in_attr_name):
        #assert out_node.has_attr(out_attr_name)
        #assert in_node.has_attr(in_attr_name)
        if out_node is in_node:
            raise ValueError("Cannot connect to same node")
        
        in_attr = in_node.attr_named(in_attr_name)
        out_attr = out_node.attr_named(out_attr_name)
        
        if in_attr.is_output():
            raise ValueError("Cannot to output attribute")
        
        if not out_attr.is_output():
            raise ValueError("Cannot connect from input attribute")
        
        if self.has_connection_to(in_node, in_attr_name):
            raise ValueError("Input attribute already has connection")
        
        if not out_attr.can_depend_on(in_attr):
            raise ValueError("Cannot connect " + str(out_attr) + " to " + str(in_attr))
        
        in_value = getattr(in_node, in_attr_name)
        out_value = getattr(out_node, out_attr_name)
        if hasattr(in_value, "get_in_connect"):
            in_attr_name = in_value.get_in_connect(out_value)
            #self.add_node(in_node)
            print 'IN ATTR NAME', in_attr_name
        key = self.create_key(out_attr_name, in_attr_name)
        self.graph.add_edge(out_node, in_node, key=key, in_attr=in_attr_name, out_attr=out_attr_name)
    
        if not nx.is_directed_acyclic_graph(self.graph):
            self.remove_edge(out_node, out_attr_name, in_node, in_attr_name)
            raise ValueError("Graph is not allowed to contain loops")
        
        if not self._custom_data.has_key('edges'):
            self._custom_data['edges'] = {}
        self._custom_data['edges'][key] = {'out_node':out_node, 
                                           'out_attr':out_attr_name,
                                           'in_node':in_node,
                                           'in_attr':in_attr_name}
        self._custom_data._p_note_change()
        cmds.save()
    
    def has_loops(self):
        if nx.topological_sort(self.graph) is None:
            return True
        found_nodes = []
        for node in self.topological_sort_iter():
            if node in found_nodes:
                return True
            found_nodes.append(node)
        return False
    
    def get_key(self, out_attr_name, in_attr_name):
        return out_attr_name + ":" + in_attr_name
    
    def create_key(self, in_attr_name, out_attr_name):
        return in_attr_name + ":" + out_attr_name
    
    def topological_sort_iter(self):
        for node in nx.topological_sort(self.graph):
            yield node
    
    def root_nodes(self):
        for node in self.graph.nodes():
            if len(self.prev(node)) == 0:
                yield node
    
    def update(self):
        for node in self.graph.nodes():
            node.prepare_update()
        root_node = list(self.topological_sort_iter())[0]
        if root_node is None: return
        self.update_node(root_node)
        self.dirty_all_nodes()
        cmds.save()
        
    def dirty_all_nodes(self):
        for node in self.nodes():
            self.dirty_node(node)
            
    def dirty_node(self, node):
        self.graph.node[node]['dirty'] = True
        
    def dirty_dependency_nodes(self, root):
        for node in nx.dfs_successor(self.graph, root).keys()[1:]:
            self.graph.node[node]['dirty'] = True
        
    def update_node(self, node):
        if not self.graph.node[node]['dirty']: 
            return
        self.pull_data(node)
    
        node.update()
        self.graph.node[node]['dirty'] = False
        for attr in node.output_attrs():
            data = getattr(node, attr.name)
            attr_edges = self.out_edges(node, attr)
            if isinstance(data, types.GeneratorType):
                for item in getattr(node, attr.name):
                    self.dirty_dependency_nodes(node)
                    for edge in self.out_edges(node, attr):
                        self.push_data(edge, item)
                    for next_node in self.next(node):
                        self.update_node(next_node)
            else:
                for next_node in self.next(node):
                    for edge in attr_edges:
                        self.push_data(edge, data)
                    self.update_node(next_node)
        
    def graphviz_layout(self):
        return nx.graphviz_layout(self.graph)
    
    def depth_node_iter(self, root=None, depth=0):
        for child in self.next(root):
            yield (depth, child)
            for depth, sub_child in self.depth_node_iter(child, depth):
                yield (depth+1, sub_child)
        
    def tree_layout(self):
        roots = list(self.root_nodes())
        tree = []
        
        pos_dict = {}
        
        tree.append(roots)
        
        for root in roots:
            
            for depth, child in self.depth_node_iter(root):
                if not len(tree) > depth+1:
                    tree.append([])
                tree[depth+1].append(child)
        
        x,y = 0,0
        xspace, yspace = 150, -40
        for depth, nodes in enumerate(tree):
            for num, node in enumerate(nodes):
                pos_dict[node] = (depth*xspace, len(nodes)*yspace - num*yspace,)
                
        return pos_dict
        
    def pull_data(self, node):
        print 'pull'
        for dirty_node in self.prev(node, dirty=True):
            self.update_node(dirty_node)
    
    def push_data(self, edge, data):
#        data = copy.deepcopy(data)
        print 'push', data, 'to', edge.in_node, edge.in_attr
        setattr(edge.in_node, edge.in_attr, data)
    
    def next(self, node, dirty=None):
        if dirty is None:
            return self.graph.successors(node)
        return [n for n in self.graph.successors(node) if self.graph.node[n]['dirty'] == dirty]
    
    def prev(self, node, dirty=None):
        if dirty is None:
            return self.graph.predecessors(node)
        return [n for n in self.graph.predecessors(node) if self.graph.node[n]['dirty'] == dirty]
                
    def out_edges(self, node, out_attr=None):
        for out_node, in_node, data in self.graph.out_edges_iter(node, data=True):
            if out_attr is None or data['out_attr'] == out_attr.name:
                yield Edge(out_node, data['out_attr'], in_node, data['in_attr'])
    
